package com.github.jiangwangyang.a2a.client;

import com.fasterxml.jackson.core.JsonProcessingException;
import io.a2a.http.A2AHttpClient;
import io.a2a.spec.*;
import io.a2a.util.Utils;
import org.springframework.http.HttpMethod;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;

import java.util.HashMap;
import java.util.Map;

import static io.a2a.util.Assert.checkNotNullParam;

public class WebClientA2AClient extends A2AClient {

    protected final Map<String, String> headers = new HashMap<>();

    public WebClientA2AClient(String agentUrl, WebClient webClient, Map<String, String> headers) {
        super(agentUrl, new WebClientA2AHttpClient(webClient));
        if (headers != null) {
            this.headers.putAll(headers);
        }
    }

    public Flux<StreamingEventKind> sendMessageFlux(MessageSendParams messageSendParams) {
        checkNotNullParam("messageSendParams", messageSendParams);
        SendStreamingMessageRequest requestBody = new SendStreamingMessageRequest.Builder()
                .jsonrpc(JSONRPCMessage.JSONRPC_VERSION)
                .method(SendStreamingMessageRequest.METHOD)
                .params(messageSendParams)
                .build();
        try {
            return ((WebClientA2AHttpClient) httpClient).requestFlux(HttpMethod.POST, agentUrl, headers, Utils.OBJECT_MAPPER.writeValueAsString(requestBody));
        } catch (JsonProcessingException e) {
            throw new RuntimeException(e);
        }
    }

    public Flux<StreamingEventKind> resubscribeToTaskFlux(TaskIdParams taskIdParams) {
        checkNotNullParam("taskIdParams", taskIdParams);
        TaskResubscriptionRequest requestBody = new TaskResubscriptionRequest.Builder()
                .jsonrpc(JSONRPCMessage.JSONRPC_VERSION)
                .method(TaskResubscriptionRequest.METHOD)
                .params(taskIdParams)
                .build();
        try {
            return ((WebClientA2AHttpClient) httpClient).requestFlux(HttpMethod.POST, agentUrl, headers, Utils.OBJECT_MAPPER.writeValueAsString(requestBody));
        } catch (JsonProcessingException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    protected A2AHttpClient.PostBuilder createPostBuilder(Object value) throws JsonProcessingException {
        A2AHttpClient.PostBuilder postBuilder = super.createPostBuilder(value);
        headers.forEach(postBuilder::addHeader);
        return postBuilder;
    }
}
